"""
DE-CN-CLIP 融合模型：将 DE-CLIP 蒸馏技术与 CN-CLIP 中文能力相结合

核心思想：
1. 基于 CN-CLIP 预训练模型，获得强大的中文理解能力
2. 添加蒸馏投影层，模拟 DE-CLIP 的知识蒸馏机制
3. 通过对比学习训练，让图像和中文文本在同一特征空间中对齐

技术原理：
- CN-CLIP：专为中文设计的视觉-语言预训练模型
- DE-CLIP：通过蒸馏提升 CLIP 性能的技术
- 融合：在 CN-CLIP 基础上添加可训练的投影层进行特征变换
"""
import torch
import torch.nn as nn
from transformers import CLIPModel, CLIPTextModel, CLIPVisionModel
from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F


class DECNClipModel(nn.Module):
    """DE-CN-CLIP 融合模型
    
    这个模型将 DE-CLIP 的蒸馏思想应用到 CN-CLIP 上：
    1. 使用预训练的 CN-CLIP 作为特征提取器
    2. 添加可学习的投影层对特征进行变换（类似蒸馏过程）
    3. 通过对比学习让图像和中文文本特征对齐
    """
    
    def __init__(self, 
                 cnclip_model_name="OFA-Sys/chinese-clip-vit-base-patch16", 
                 temperature=0.07):
        """初始化模型
        
        Args:
            cnclip_model_name: CN-CLIP 预训练模型名称
            temperature: 对比学习的温度参数，控制相似度计算的平滑程度
        """
        super().__init__()
        
        # === 加载 CN-CLIP 预训练模型 ===
        # CN-CLIP 是专为中文优化的 CLIP 模型，能更好理解中文语义
        self.cnclip = AutoModel.from_pretrained(cnclip_model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(cnclip_model_name)
        
        # === 蒸馏投影层（模拟 DE-CLIP 的知识蒸馏） ===
        # DE-CLIP 通过添加投影层来改进特征表示，这里借鉴这个思想
        self.text_projection = nn.Linear(512, 512)    # 中文文本特征投影
        self.visual_projection = nn.Linear(512, 512)  # 视觉特征投影
        
        # === 训练超参数 ===
        self.temperature = temperature  # 温度参数：越小相似度差异越大
        self.dropout = nn.Dropout(0.1)  # 防止过拟合
        
    def encode_text(self, text_inputs):
        """编码中文文本为特征向量
        
        步骤：
        1. 使用 CN-CLIP 的文本编码器提取原始特征
        2. 通过投影层进行特征变换（蒸馏步骤）
        3. 应用 dropout 防止过拟合
        4. L2 标准化确保特征在单位球面上
        
        Args:
            text_inputs: tokenized 的中文文本输入
        Returns:
            标准化后的文本特征向量 [batch_size, 512]
        """
        # 第1步：CN-CLIP 文本编码器提取特征
        text_features = self.cnclip.get_text_features(**text_inputs)
        
        # 第2步：蒸馏投影变换（核心改进点）
        # 这个投影层学习如何改进原始 CN-CLIP 特征
        projected_features = self.text_projection(text_features)
        projected_features = self.dropout(projected_features)
        
        # 第3步：L2标准化，确保所有特征向量长度为1
        # 这样可以用余弦相似度计算，更稳定
        return F.normalize(projected_features, dim=-1)
    
    def encode_image(self, image_inputs):
        """编码图像为特征向量
        
        步骤与文本编码类似：
        1. CN-CLIP 视觉编码器提取图像特征
        2. 投影层进行特征变换
        3. 标准化输出
        
        Args:
            image_inputs: 预处理后的图像张量
        Returns:
            标准化后的图像特征向量 [batch_size, 512]
        """
        # 使用 CN-CLIP 的视觉编码器
        image_features = self.cnclip.get_image_features(image_inputs)
        
        # 蒸馏投影变换
        projected_features = self.visual_projection(image_features)
        projected_features = self.dropout(projected_features)
        
        # L2标准化
        return F.normalize(projected_features, dim=-1)
    
    def forward(self, image_inputs, text_inputs):
        """模型前向传播
        
        核心逻辑：
        1. 分别编码图像和文本得到特征向量
        2. 计算所有图像-文本对的相似度矩阵
        3. 用温度参数调节相似度分布
        
        这就是对比学习的核心：让匹配的图文对相似度高，不匹配的低
        
        Args:
            image_inputs: 图像张量 [batch_size, 3, 224, 224]
            text_inputs: 文本输入字典 {input_ids, attention_mask}
            
        Returns:
            包含特征和相似度的字典
        """
        # 编码图像和文本
        image_features = self.encode_image(image_inputs)  # [batch_size, 512]
        text_features = self.encode_text(text_inputs)     # [batch_size, 512]
        
        # 计算相似度矩阵：每个图像与每个文本的相似度
        # 矩阵乘法：[batch_size, 512] × [512, batch_size] = [batch_size, batch_size]
        logits = torch.matmul(image_features, text_features.T) / self.temperature
        
        # logits[i,j] 表示第i张图像与第j个文本的相似度
        # 对角线元素 logits[i,i] 是匹配对，应该最大
        
        return {
            'image_features': image_features,
            'text_features': text_features, 
            'logits': logits  # 相似度矩阵
        }
    
    def compute_loss(self, image_inputs, text_inputs):
        """计算对比学习损失
        
        对比学习的目标：
        - 让匹配的图文对相似度尽可能高
        - 让不匹配的图文对相似度尽可能低
        
        具体做法：
        1. 图像-文本检索：给定图像，正确文本应该排第一
        2. 文本-图像检索：给定文本，正确图像应该排第一
        3. 两个方向的损失取平均
        
        Args:
            image_inputs: 图像输入
            text_inputs: 文本输入
            
        Returns:
            对比学习损失值
        """
        outputs = self(image_inputs, text_inputs)
        logits = outputs['logits']  # [batch_size, batch_size]
        
        # 创建标签：第i个样本的正确匹配是第i个
        batch_size = logits.shape[0]
        labels = torch.arange(batch_size, device=logits.device)  # [0,1,2,...,batch_size-1]
        
        # === 双向对比损失 ===
        # 图像到文本：每行是一张图像对所有文本的相似度
        # 目标：logits[i, i] 应该是第i行的最大值
        loss_i2t = F.cross_entropy(logits, labels)
        
        # 文本到图像：每列是一个文本对所有图像的相似度  
        # 目标：logits[i, i] 应该是第i列的最大值
        loss_t2i = F.cross_entropy(logits.T, labels)
        
        # 对称损失：两个方向平均
        return (loss_i2t + loss_t2i) / 2


def create_decnclip_model(model_name="OFA-Sys/chinese-clip-vit-base-patch16"):
    """创建 DE-CN-CLIP 模型的工厂函数
    
    Args:
        model_name: 使用的 CN-CLIP 预训练模型名称
        
    Returns:
        初始化好的 DECNClipModel 实例
    """
    return DECNClipModel(model_name)


if __name__ == "__main__":
    # 测试模型创建
    print("正在创建 DE-CN-CLIP 模型...")
    model = create_decnclip_model()
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"模型创建成功！")
    print(f"总参数数量: {total_params:,}")
    print(f"可训练参数: {trainable_params:,}")
    print(f"模型结构预览:")
    print(f"- CN-CLIP 预训练模型")
    print(f"- 文本投影层: 512 -> 512") 
    print(f"- 视觉投影层: 512 -> 512")
    print(f"- Dropout 层: p=0.1")